import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from scipy import ndimage


def generate_blue_noise_mask(size):
    """生成蓝噪声掩模（基于void-and-cluster方法简化版）"""
    # 步骤1：初始随机二进制模式
    mask = np.random.rand(size, size) > 0.5

    # 步骤2：通过滤波优化为蓝噪声特性
    kernel = np.array([[0.05, 0.2, 0.05],
                       [0.2, -1.0, 0.2],
                       [0.05, 0.2, 0.05]])
    for _ in range(10):  # 迭代优化次数
        conv = ndimage.convolve(mask.astype(float), kernel, mode='wrap')
        mask = (conv < 0)  # 更新掩模
    return mask.astype(float)


def dart_blue_noise_sampling(image, sample_ratio=0.3):
    """DART蓝噪声采样"""
    # 转换为灰度图
    if image.ndim == 3:
        gray = np.mean(image, axis=2).astype(float)
    else:
        gray = image.astype(float)

    # 生成蓝噪声掩模
    h, w = gray.shape
    blue_noise = generate_blue_noise_mask(max(h, w))[:h, :w]

    # 归一化并采样
    gray_normalized = (gray - gray.min()) / (gray.max() - gray.min())
    sampled = (gray_normalized > blue_noise * sample_ratio).astype(np.uint8) * 255

    return sampled


# 读取图片
image = np.array(Image.open("face.png").convert("L"))  # 转为灰度图

# 执行蓝噪声采样
sampled = dart_blue_noise_sampling(image, sample_ratio=0.2)

# 显示结果
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.imshow(image, cmap='gray')
plt.title("原图")
plt.axis('off')

plt.subplot(1, 2, 2)
plt.imshow(sampled, cmap='gray')
plt.title("蓝噪声采样结果")
plt.axis('off')
plt.show()

# 保存结果
Image.fromarray(sampled).save("blue_noise_sampled.png")